import redis
import re

def init_unassigned_annos(redis_client, task_id=None):
    """
    查找status为0的ann_id并添加到unassigned_annos集合
    
    :param redis_client: Redis客户端实例
    :param task_id: 可选，指定任务ID，None表示所有任务
    :return: 成功添加的ann_id数量
    """
    # 构建键匹配模式
    if task_id:
        pattern = f"task_info:{task_id}:anno:*"
    else:
        pattern = "task_info:*:anno:*"
    
    # 正则表达式用于提取ann_id
    # 匹配格式: task_info:{task_id}:anno:{ann_id}
    regex = re.compile(r"task_info:(\w+):anno:(\w+)")
    
    cursor = 0
    count = 0
    pipeline = redis_client.pipeline()  # 使用管道批量操作
    
    try:
        while True:
            # 扫描匹配的键
            cursor, keys = redis_client.scan(
                cursor=cursor,
                match=pattern,
                count=1000  # 每次扫描尝试返回1000个键
            )
            
            if not keys:
                if cursor == 0:
                    break
                continue
            
            # 批量获取所有键的status字段
            for key in keys:
                pipeline.hget(key, "status")
            
            # 执行批量查询
            statuses = pipeline.execute()
            
            # 处理结果
            for i, key in enumerate(keys):
                status = statuses[i]
                
                # 检查status是否为0
                if status == b'0':  # Redis返回的是字节类型
                    # 提取ann_id
                    match = regex.match(key.decode())
                    if match:
                        task_id = match.group(1)
                        ann_id = match.group(2)
                        # 添加到集合
                        redis_client.sadd(f"task_info:{task_id}:unassigned_annos", ann_id)
                        count += 1
            
            if cursor == 0:
                break
                
    except Exception as e:
        print(f"处理过程中发生错误: {str(e)}")
        raise
    
    return count

if __name__ == "__main__":
    # 连接Redis
    r = redis.Redis(
        host='localhost',
        port=6379,
        db=0,password="123456",
        # decode_responses=False  # 保持字节类型以统一处理
    )
    
    # 示例：初始化所有任务中status为0的ann_id
    total = init_unassigned_annos(r)
    print(f"成功将{total}个未分配的ann_id添加到unassigned_annos集合")
    
    # 示例：初始化指定任务ID的未分配ann_id
    # task_id = 3
    # total = init_unassigned_annos(r, task_id)
    # print(f"成功将任务{task_id}中{total}个未分配的ann_id添加到unassigned_annos集合")
